#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Software: PyCharm
# @Version : Python-
# @Author  : Shengji He
# @Email   : hsjbit@163.com
# @File    : CountRangeSum.py
# @Time    : 2020/11/7 13:04
# @Description:
from typing import List
import bisect


class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        """
        Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
        Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

        Note:
            A naive algorithm of O(n2) is trivial. You MUST do better than that.

        Example:
            Input: nums = [-2,5,-1], lower = -2, upper = 2,
            Output: 3
            Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.

        Constraints:
            - 0 <= nums.length <= 10^4

        :param nums:
        :param lower:
        :param upper:
        :return:
        """
        # n = len(nums)
        # dp = [[0 for _ in range(n)] for _ in range(n)]
        # ans = []
        # for i in range(n):
        #     dp[i][i] = nums[i]
        #     if lower <= nums[i] <= upper:
        #         ans.append([i, i])
        # row = 0
        # while row < n - 1:
        #     column = row + 1
        #     for i in range(n - column):
        #         dp[row][column + i] = dp[row][column + i - 1] + dp[column + i][column + i]
        #         if lower <= dp[row][column + i] <= upper:
        #             ans.append([row, column + i])
        #     row += 1
        # return len(ans)
        res = cur = 0
        s = [0]
        for v in nums:
            cur += v
            res += bisect.bisect_right(s, cur - lower) - bisect.bisect_left(s, cur - upper)
            bisect.insort_right(s, cur)

        return res


if __name__ == '__main__':
    S = Solution()
    nums = [-2, 5, -1]
    lower = -2
    upper = 2
    print(S.countRangeSum(nums, lower, upper))
    print('done')
